package ipaX

import (
	"fmt"
	"runtime"

	"github.com/crate-crypto/go-ipa/bandersnatch"
	"github.com/crate-crypto/go-ipa/bandersnatch/fr"
	"github.com/crate-crypto/go-ipa/common"
	"github.com/crate-crypto/go-ipa/ipa"
	gthCrypto "github.com/ethereum/go-ethereum/crypto"
)

// type IPAConf struct {
// 	*ipa.IPAConfig
// }

//
// type PrecomputedWeights struct {
// 		This stores A'(x_i) and 1/A'(x_i)
// A'(x_i) = (x_i-j)!-> j(0, 256) i->(0, 256)
// [A'(p_1), A'(p_2), ..., A'(p_256), 1/A'(p_1), ..., 1/A'(p_256)]
//	barycentricWeights []fr.Element
// 	This stores 1/k and -1/k , for k in [1, 255]
//  [1/1, 1/2, ''', 1/255, -1/1, ..., -1/255]
//	invertedDomain []fr.Element
// }

// []byte to element
func genRoot(ins [][]byte) []fr.Element {
	n := len(ins)
	if len(ins) > 256 {
		panic("polynomial cannot exceed 256 coefficients")
	}

	root := make([]fr.Element, 256, 256)
	for i, one := range ins {
		hash := gthCrypto.Keccak256(one)
		root[i].SetBytes(hash)
	}

	for i := n; i < 256; i++ {
		root[i] = fr.Zero()
	}
	return root
}

// 利用插值法，生成多项式
func polyGen(root []fr.Element) []fr.Element {
	var negFirst fr.Element
	negFirst.Neg(&root[0])
	poly := []fr.Element{negFirst, fr.One()} // (x - root[0]) => [-root[0], 1]
	for i := 1; i < len(root); i++ {
		var tmp fr.Element
		tmp.Neg(&root[i])
		poly = polynomialMul(poly, []fr.Element{tmp, fr.One()})
	}
	return poly
}

// c = a0*g0 + a1g1 + ..... a13*g13
func CommitPoly(group_elements []bandersnatch.PointAffine, polynomial []fr.Element) bandersnatch.PointAffine {
	if len(group_elements) != len(polynomial) {
		panic(fmt.Sprintf("diff sizes, %d != %d", len(group_elements), len(polynomial)))
	}

	var result bandersnatch.PointProj
	result.Identity()

	var res, err = result.MultiExp(group_elements, polynomial, bandersnatch.MultiExpConfig{NbTasks: runtime.NumCPU(), ScalarsMont: true})
	if err != nil {
		panic("mult exponentiation was not successful. TODO: replace panics by bubbling up error")
	}

	var resAffine bandersnatch.PointAffine
	resAffine.FromProj(res)
	return resAffine
}

// commitment = [vector a] * [vector g]
// C = a> * g>  +  a> * b> * q
func CreateIPAProof(transcript *common.Transcript, ic *ipa.IPAConfig, commitment bandersnatch.PointAffine, a []fr.Element, eval_point fr.Element) ipa.IPAProof {
	transcript.DomainSep("ipa")
	// vector b = b> = 1, e, e^2, e^3
	b := ic.PrecomputedWeights.ComputeBarycentricCoefficients(eval_point)
	inner_prod := ipa.InnerProd(a, b)

	transcript.AppendPoint(&commitment, "C")
	transcript.AppendScalar(&eval_point, "input point")
	transcript.AppendScalar(&inner_prod, "output point")
	w := transcript.ChallengeScalar("w")

	var q bandersnatch.PointAffine
	q.ScalarMul(&ic.SRSPrecompPoints.Q, &w)

	num_rounds := compute_num_rounds(common.POLY_DEGREE)
	current_basis := ic.SRSPrecompPoints.SRS

	L := make([]bandersnatch.PointAffine, num_rounds)
	R := make([]bandersnatch.PointAffine, num_rounds)

	for i := 0; i < int(num_rounds); i++ {

		a_L, a_R := splitScalars(a)

		b_L, b_R := splitScalars(b)

		G_L, G_R := splitPoints(current_basis)

		z_L := ipa.InnerProd(a_R, b_L)
		z_R := ipa.InnerProd(a_L, b_R)

		C_L_1 := CommitPoly(G_L, a_R)
		C_L := CommitPoly([]bandersnatch.PointAffine{C_L_1, q}, []fr.Element{fr.One(), z_L})

		C_R_1 := CommitPoly(G_R, a_L)
		C_R := CommitPoly([]bandersnatch.PointAffine{C_R_1, q}, []fr.Element{fr.One(), z_R})

		L[i] = C_L
		R[i] = C_R

		transcript.AppendPoint(&C_L, "L")
		transcript.AppendPoint(&C_R, "R")
		x := transcript.ChallengeScalar("x")

		var xInv fr.Element
		xInv.Inverse(&x)

		// TODO: We could use a for loop here like in the Rust code
		a = foldScalars(a_L, a_R, x)
		b = foldScalars(b_L, b_R, xInv)

		current_basis = foldPoints(G_L, G_R, xInv)
	}

	if len(a) != 1 {
		panic("length of `a` should be 1 at the end of the reduction")
	}

	return ipa.IPAProof{
		L:        L,
		R:        R,
		A_scalar: a[0],
	}

}

func CheckIPAProof(transcript *common.Transcript, ic *ipa.IPAConfig, commitment bandersnatch.PointAffine, proof ipa.IPAProof, eval_point fr.Element, inner_prod fr.Element) bool {
	transcript.DomainSep("ipa")

	if len(proof.L) != len(proof.R) {
		panic("L and R should be the same size")
	}
	if len(proof.L) != int(compute_num_rounds(common.POLY_DEGREE)) {
		panic("The number of points for L or R should be equal to the number of rounds")
	}

	b := ic.PrecomputedWeights.ComputeBarycentricCoefficients(eval_point)

	transcript.AppendPoint(&commitment, "C")
	transcript.AppendScalar(&eval_point, "input point")
	transcript.AppendScalar(&inner_prod, "output point")

	w := transcript.ChallengeScalar("w")

	var q bandersnatch.PointAffine
	q.ScalarMul(&ic.SRSPrecompPoints.Q, &w)

	var qy bandersnatch.PointAffine
	qy.ScalarMul(&q, &inner_prod)
	commitment.Add(&commitment, &qy)

	challenges := generateChallenges(transcript, &proof)
	challenges_inv := make([]fr.Element, len(challenges))

	// Compute expected commitment
	for i := 0; i < len(challenges); i++ {
		x := challenges[i]
		L := proof.L[i]
		R := proof.R[i]

		var xInv fr.Element
		xInv.Inverse(&x)

		challenges_inv[i] = xInv
		// C' = x*C_l + C + 1/x*C_r
		commitment = CommitPoly([]bandersnatch.PointAffine{commitment, L, R}, []fr.Element{fr.One(), x, xInv})
	}

	current_basis := ic.SRSPrecompPoints.SRS

	for i := 0; i < len(challenges); i++ {

		G_L, G_R := splitPoints(current_basis)

		b_L, b_R := splitScalars(b)

		xInv := challenges_inv[i]
		// b' = b_l + 1/x * b_r
		b = foldScalars(b_L, b_R, xInv)
		// g' = g_l + 1/x*g_r
		current_basis = foldPoints(G_L, G_R, xInv)
	}

	if len(b) != len(current_basis) {
		panic("reduction was not correct")
	}

	if len(b) != 1 {
		panic("`b` and `G` should be 1")
	}

	b0 := b[0]

	var got bandersnatch.PointAffine
	//  G[0] * a + (a * b) * Q;
	var part_1 bandersnatch.PointAffine
	part_1.ScalarMul(&current_basis[0], &proof.A_scalar)

	var part_2 bandersnatch.PointAffine
	var part_2a fr.Element

	part_2a.Mul(&b0, &proof.A_scalar)
	part_2.ScalarMul(&q, &part_2a)
	// ag + abq = C
	got.Add(&part_1, &part_2)

	return got.Equal(&commitment)
}
